package com.future;

import com.future.util.TreeNode;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
 * Description: 863. 二叉树中所有距离为 K 的结点
 * <p>
 * 输入：root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, k = 2
 * 输出：[7,4,1]
 * 解释：所求结点为与目标结点（值为 5）距离为 2 的结点，值分别为 7，4，以及 1
 * <p>
 * 来源：力扣（LeetCode）
 * 链接：https://leetcode.cn/problems/all-nodes-distance-k-in-binary-tree
 * 著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。
 *
 * @author weiruibai.vendor
 * Date: 2023/1/3 10:53
 */
public class Solution_863 {

    private static Solution_863 instance = new Solution_863();

    public static void main(String[] args) {
        TreeNode root = new TreeNode(3);
        TreeNode target = new TreeNode(5);
        root.left = target;
        root.left.left = new TreeNode(6);
        TreeNode node = new TreeNode(2);
        root.left.right = node;
        node.left = new TreeNode(7);
        node.right = new TreeNode(4);

        root.right = new TreeNode(1);
        root.right.left = new TreeNode(0);
        root.right.right = new TreeNode(8);

        int k = 2;
        List<Integer> list = instance.distanceK(root, target, k);

        System.out.println(list);
    }

    public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
        Map<TreeNode, TreeNode> childParent = new HashMap<>();
        childParent.put(root, null);
        List<Integer> ans = new ArrayList<>();
        buildMap(root, childParent);
        Set<TreeNode> repeat = new HashSet<>();
        getDistanceK(target, childParent, ans, repeat, k);
        return ans;
    }

    private void getDistanceK(TreeNode target, Map<TreeNode, TreeNode> childParent, List<Integer> ans, Set<TreeNode> repeat, int k) {
        if (target == null || repeat.contains(target)) {
            return;
        }
        if (k == 0) {
            repeat.add(target);
            ans.add(target.val);
            return;
        }
        k--;
        repeat.add(target);
        // 向父节点追溯
        TreeNode curNode = childParent.get(target);
        if (curNode == null) {
            // 来到根节点
            getDistanceK(target.left, childParent, ans, repeat, k);
            getDistanceK(target.right, childParent, ans, repeat, k);
        } else {
            // 向根节点追溯
            getDistanceK(curNode, childParent, ans, repeat, k);
        }
        getDistanceK(target.left, childParent, ans, repeat, k);
        getDistanceK(target.right, childParent, ans, repeat, k);
        repeat.remove(target);
    }

    private void buildMap(TreeNode root, Map<TreeNode, TreeNode> childParent) {
        if (root == null) {
            return;
        }
        childParent.put(root.left, root);
        buildMap(root.left, childParent);
        childParent.put(root.right, root);
        buildMap(root.right, childParent);
    }
}
